// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <deque>
#include "kllm/core/metric.h"
#include <mutex>
#include <condition_variable>

namespace kllm {


    struct server_queue {
        int id = 0;
        bool running;

        // queues
        std::deque<server_task> queue_tasks;
        std::deque<server_task> queue_tasks_deferred;

        std::mutex mutex_tasks;
        std::condition_variable condition_tasks;

        // callback functions
        std::function<void(server_task)> callback_new_task;
        std::function<void(void)>        callback_update_slots;

        // Add a new task to the end of the queue
        int post(server_task task, bool front = false) {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            if (task.id == -1) {
                task.id = id++;
            }
            QUE_DBG("new task, id = %d, front = %d\n", task.id, front);
            if (front) {
                queue_tasks.push_front(std::move(task));
            } else {
                queue_tasks.push_back(std::move(task));
            }
            condition_tasks.notify_one();
            return task.id;
        }

        // multi-task version of post()
        int post(std::vector<server_task> & tasks, bool front = false) {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            for (auto & task : tasks) {
                if (task.id == -1) {
                    task.id = id++;
                }
                QUE_DBG("new task, id = %d/%d, front = %d\n", task.id, (int) tasks.size(), front);
                if (front) {
                    queue_tasks.push_front(std::move(task));
                } else {
                    queue_tasks.push_back(std::move(task));
                }
            }
            condition_tasks.notify_one();
            return 0;
        }

        // Add a new task, but defer until one slot is available
        void defer(server_task task) {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            QUE_DBG("defer task, id = %d\n", task.id);
            queue_tasks_deferred.push_back(std::move(task));
            condition_tasks.notify_one();
        }

        // Get the next id for creating a new task
        int get_new_id() {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            int new_id = id++;
            return new_id;
        }

        // Register function to process a new task
        void on_new_task(std::function<void(server_task)> callback) {
            callback_new_task = std::move(callback);
        }

        // Register the function to be called when all slots data is ready to be processed
        void on_update_slots(std::function<void(void)> callback) {
            callback_update_slots = std::move(callback);
        }

        // Call when the state of one slot is changed, it will move one task from deferred to main queue
        void pop_deferred_task() {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            if (!queue_tasks_deferred.empty()) {
                queue_tasks.emplace_back(std::move(queue_tasks_deferred.front()));
                queue_tasks_deferred.pop_front();
            }
            condition_tasks.notify_one();
        }

        // end the start_loop routine
        void terminate() {
            std::unique_lock<std::mutex> lock(mutex_tasks);
            running = false;
            condition_tasks.notify_all();
        }

        /**
         * Main loop consists of these steps:
         * - Wait until a new task arrives
         * - Process the task (i.e. maybe copy data into slot)
         * - Check if multitask is finished
         * - Update all slots
         */
        void start_loop() {
            running = true;

            while (true) {
                QUE_DBG("%s", "processing new tasks\n");

                while (true) {
                    std::unique_lock<std::mutex> lock(mutex_tasks);
                    if (queue_tasks.empty()) {
                        lock.unlock();
                        break;
                    }
                    server_task task = queue_tasks.front();
                    queue_tasks.pop_front();
                    lock.unlock();

                    QUE_DBG("processing task, id = %d\n", task.id);
                    callback_new_task(std::move(task));
                }

                // all tasks in the current loop is processed, slots data is now ready
                QUE_DBG("%s", "update slots\n");

                callback_update_slots();

                QUE_DBG("%s", "waiting for new tasks\n");
                {
                    std::unique_lock<std::mutex> lock(mutex_tasks);
                    if (queue_tasks.empty()) {
                        if (!running) {
                            QUE_DBG("%s", "terminate\n");
                            return;
                        }
                        condition_tasks.wait(lock, [&]{
                            return (!queue_tasks.empty() || !running);
                        });
                    }
                }
            }
        }
    };

}  // namespace kllm
